import os
import fire
import sys
import time

# Add the project root to the Python path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..'))
sys.path.append(project_root)

from omics_bert_label_encoding.encoding.read_encoding_embeddings import EncodingEmbeddingReader

def main(dataset_path: str = os.path.join(project_root, "embeddings_local/charls/v1/encoding_embeddings/clhls_2000_2018/")):
    """Browse embeddings in a dataset."""
    with EncodingEmbeddingReader(dataset_path) as reader:
        print(f"Browsing embeddings in {dataset_path}")

        # 统计embedding有多少个
        print(f"Total embeddings: {reader.count_embeddings()}")

        # 示例：取出key为10000对应的embedding
        embedding_1 = reader.get_encoding_embedding(1)
        print(f"Embedding for key 1: {embedding_1}")

        # 示例：取出[1, 2]对应的embedding
        embeddings_2 = reader.get_encodings_embeddings_batch([1, 2])
        print(f"Embeddings for keys [1, 2]: {embeddings_2}")

        # 打印前100个embedding，每个间隔0.5秒
        print("Printing first 100 embeddings with 0.5s interval...")

        for i, (key, embedding) in enumerate(reader.iterate_embeddings()):
            if i >= 100:
                print("\nReached 100 embeddings, stopping.")
                break

            print(f"\n--- Embedding {i+1} ---")
            print(f"Key: {key}")
            print(f"Embedding: {embedding}")
            time.sleep(0.5)

if __name__ == '__main__':
    fire.Fire(main)